import React, { useState } from 'react';
import { Tree } from 'antd';

const { TreeNode } = Tree;
/**
 * @description: 受控树
 * @param [{}]
 * @return: React Node
 */
const TreeControlled = function TreeControlled(props) {
  const {
    data = [],
    checkable = false,
    selectable = false,
    checkedKeys,
    setCheckedKeys,
    expandedKeys,
    setExpandedKeys,
    selectedKeys,
    setSelectedKeys,
    setHalfCheckedKeys,
    renderTreeNodes,
    renderTitle = ({ title }) => title,
    autoExpandParent: autoExpandParentProp = true,
    ...rest
  } = props;

  const _renderTreeNodes = data =>
    data.map(item => {
      if (item.children) {
        return (
          <TreeNode key={item.key} {...item} title={renderTitle(item)} dataRef={item}>
            {_renderTreeNodes(item.children)}
          </TreeNode>
        );
      }
      return <TreeNode key={item.key} {...item} title={renderTitle(item)} />;
    });

  const [_checkedKeys, _setCheckedKeys] = useState([]);
  const [_expandedKeys, _setExpandedKeys] = useState([]);
  const [_selectedKeys, _setSelectedKeys] = useState([]);
  const [autoExpandParent, setAutoExpandParent] = useState(autoExpandParentProp);

  const _onExpand = keys => {
    setAutoExpandParent(false);
    if (setExpandedKeys) {
      setExpandedKeys(keys);
    } else {
      _setExpandedKeys(keys);
    }
  };

  const _onCheck = (keys, { halfCheckedKeys }) => {
    setAutoExpandParent(false);
    if (setCheckedKeys) {
      setCheckedKeys(keys);
    } else {
      _setCheckedKeys(keys);
    }
    if (setHalfCheckedKeys) {
      setHalfCheckedKeys(halfCheckedKeys);
    }
  };

  return (
    <Tree
      checkable={checkable}//复选框
      selectable={selectable}//是否可选中
      autoExpandParent={autoExpandParent}//是否自动展开父节点
      checkedKeys={checkedKeys || _checkedKeys}
      onCheck={_onCheck}//点击复选框触发
      onExpand={_onExpand}//展开/收起节点时触发
      expandedKeys={expandedKeys || _expandedKeys}//（受控）展开指定的树节点
      selectedKeys={selectedKeys || _selectedKeys}//（受控）设置选中的树节点
      onSelect={setSelectedKeys || _setSelectedKeys}//点击树节点触发
      {...rest}
    >
      {renderTreeNodes ? renderTreeNodes(data) : _renderTreeNodes(data)}
    </Tree>
  );
};

export default TreeControlled;
